# -*- coding: UTF-8 -*-
"""
@File    :   convert_ccpd_ori_fmt.py
@Author  :   AnimateX
@Contact :   animatex@163.com
@License :   Copyright © 2022-2023 RockChips. All rights reserved.

@Create Time           @Software Version
-----------------      -----------------
2023/12/4 16:50        Version_1.0   

------------------------------------------------------------------
@Description :

------------------------------------------------------------------
"""
import os
import json
import shutil
import inspect
from tqdm import tqdm
from natsort import natsorted


class DataProcess:
    def __init__(self):
        pass

    @staticmethod
    def get_debug_info():
        frame = inspect.currentframe()
        caller = frame.f_back

        # 获取调用者的函数名
        function_name = caller.f_code.co_name

        # 获取调用者所在的文件名
        file_name = caller.f_code.co_img_name

        # 获取行号
        line_num = caller.f_lineno

        return file_name, function_name, line_num

    @staticmethod
    def find_jpg_files_scandir(directory):
        jpg_files = []
        for entry in os.scandir(directory):
            if entry.is_dir():
                jpg_files.extend(DataProcess().find_jpg_files_scandir(entry.path))
            elif entry.is_file() and entry.name.endswith('.jpg'):
                jpg_files.append(entry.path)

        return jpg_files

    @staticmethod
    def get_car_license_info(img_path: str):
        if not os.path.exists(img_path):
            file_name, func_name, line_num = DataProcess().get_debug_info()
            print(f' [ERR] {func_name}: {line_num}, img_path not exists!')
            return

        img_name = img_path.split('/')[-1]
        # 分割文件名以获取不同的字段
        parts = img_name.split('-')
        if len(parts) < 6:
            return {}

        area = parts[0]
        tilt_degrees = parts[1]
        h_tilt_degree = int(tilt_degrees.split('_')[0])
        v_tilt_degree = int(tilt_degrees.split('_')[1])
        bbox_coords = parts[2]
        bbox_lu_x = int(bbox_coords.split('_')[0].split('&')[0])
        bbox_lu_y = int(bbox_coords.split('_')[0].split('&')[1])
        bbox_rd_x = int(bbox_coords.split('_')[1].split('&')[0])
        bbox_rd_y = int(bbox_coords.split('_')[1].split('&')[1])
        vertices = parts[3].split('_')
        vertices_rd_x = int(vertices[0].split('&')[0])
        vertices_rd_y = int(vertices[0].split('&')[1])
        vertices_ld_x = int(vertices[1].split('&')[0])
        vertices_ld_y = int(vertices[1].split('&')[1])
        vertices_lu_x = int(vertices[2].split('&')[0])
        vertices_lu_y = int(vertices[2].split('&')[1])
        vertices_ru_x = int(vertices[3].split('&')[0])
        vertices_ru_y = int(vertices[3].split('&')[1])
        plate_numbers = parts[4].split('_')
        brightness = parts[5]
        blur = parts[6].split('.')[0]  # 移除文件扩展名

        # 定义省份、字母和数字的数组, 'O' 指无效信息，比如无车牌
        provinces = ["皖", "沪", "津", "渝", "冀", "晋",
                     "蒙", "辽", "吉", "黑", "苏", "浙",
                     "京", "闽", "赣", "鲁", "豫", "鄂",
                     "湘", "粤", "桂", "琼", "川", "贵",
                     "云", "藏", "陕", "甘", "青", "宁",
                     "新", "警", "学", "O"]

        alphabets = ['A', 'B', 'C', 'D', 'E', 'F', 'G',
                     'H', 'J', 'K', 'L', 'M', 'N', 'P',
                     'Q', 'R', 'S', 'T', 'U', 'V', 'W',
                     'X', 'Y', 'Z', 'O']

        ads = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H',
               'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R',
               'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
               '0', '1', '2', '3', '4', '5', '6', '7',
               '8', '9', 'O']

        # 解析车牌号码
        province = provinces[int(plate_numbers[0])]
        alphabet = alphabets[int(plate_numbers[1])]
        ad = ''.join([ads[int(num)] for num in plate_numbers[2:]])

        # 创建字典来存储解析后的信息
        parsed_info = {
            'area': area,
            'tilt_degrees': {
                'h_degree': h_tilt_degree,
                'v_degree': v_tilt_degree,
            },
            'bbox_coords': {
                'bbox_lu': {
                    'x': bbox_lu_x,
                    'y': bbox_lu_y,
                },
                'bbox_rd': {
                    'x': bbox_rd_x,
                    'y': bbox_rd_y,
                },
            },
            'vertices': {
                'vertices_lu': {
                    'x': vertices_lu_x,
                    'y': vertices_lu_y,
                },
                'vertices_ru': {
                    'x': vertices_ru_x,
                    'y': vertices_ru_y,
                },
                'vertices_ld': {
                    'x': vertices_ld_x,
                    'y': vertices_ld_y,
                },
                'vertices_rd': {
                    'x': vertices_rd_x,
                    'y': vertices_rd_y,
                },
            },
            'license_plate': {
                'province': province,
                'alphabet': alphabet,
                'ad': ad,
            },
            'brightness': brightness,
            'blur': blur,
        }

        return parsed_info

    @staticmethod
    def convert_info_to_yolo_format(img_info_dict: dict, img_w=720, img_h=1160):
        bbox_coords = img_info_dict['bbox_coords']
        s_x = bbox_coords['bbox_lu']['x']
        s_y = bbox_coords['bbox_lu']['y']
        e_x = bbox_coords['bbox_rd']['x']
        e_y = bbox_coords['bbox_rd']['y']

        roi_w = e_x - s_x
        roi_h = e_y - s_y
        roi_c_x = (s_x + e_x) / 2.0
        roi_c_y = (s_y + e_y) / 2.0

        # norm
        c_x = round(roi_c_x / img_w, 6)
        c_y = round(roi_c_y / img_h, 6)
        roi_w = round(roi_w / img_w, 6)
        roi_h = round(roi_h / img_h, 6)

        return c_x, c_y, roi_w, roi_h

    @staticmethod
    def create_dst_dataset_path(dst_dataset_dir_path: str):
        dst_images_path = os.path.join(dst_dataset_dir_path, 'images')
        dst_images_train_path = os.path.join(dst_images_path, 'train')
        dst_images_valid_path = os.path.join(dst_images_path, 'val')
        dst_labels_path = os.path.join(dst_dataset_dir_path, 'labels')
        dst_labels_train_path = os.path.join(dst_labels_path, 'train')
        dst_labels_valid_path = os.path.join(dst_labels_path, 'val')

        if not os.path.exists(dst_images_train_path):
            os.makedirs(dst_images_train_path)

        if not os.path.exists(dst_images_valid_path):
            os.makedirs(dst_images_valid_path)

        if not os.path.exists(dst_labels_train_path):
            os.makedirs(dst_labels_train_path)

        if not os.path.exists(dst_labels_valid_path):
            os.makedirs(dst_labels_valid_path)

        res_path_dict = {
            'train_images_path': dst_images_train_path,
            'valid_images_path': dst_images_valid_path,
            'train_labels_path': dst_labels_train_path,
            'valid_labels_path': dst_labels_valid_path
        }

        return res_path_dict

    @staticmethod
    def get_all_ccpd_img_list(ccpd_dataset_dir_path: str, dst_dataset_path: str, data_category='ev', cate='0', save_info_name='all_info.json', val_ratio=0.9):
        if not os.path.exists(ccpd_dataset_dir_path):
            file_name, func_name, line_num = YoloFmtDataset().get_debug_info()
            print(f' [ERR] {func_name}: {line_num}, ccpd dataset directory not exists!')
            return

        # safe check.
        res_path_dict = DataProcess().create_dst_dataset_path(dst_dataset_path)

        # Get save info path
        save_info_path = os.path.join(dst_dataset_path, save_info_name)

        # Get all image path list.
        img_path_list = DataProcess().find_jpg_files_scandir(ccpd_dataset_dir_path)

        # Get all img num.
        all_img_cnt = len(img_path_list)
        train_img_cnt = int(all_img_cnt * val_ratio)
        # valid_img_cnt = all_img_cnt - train_img_cnt

        # Iter.
        for idx, img_path in enumerate(tqdm(img_path_list, position=0, desc=' Processing CCPD data...', colour='cyan')):
            img_info_dict = DataProcess().get_car_license_info(img_path)
            if len(img_info_dict) == 0:
                continue

            # Get car license roi position.
            c_x, c_y, w, h = DataProcess().convert_info_to_yolo_format(img_info_dict)

            # save yolo train labels
            label_info = f'{cate} {c_x} {c_y} {w} {h}'

            new_image_name = f'{data_category}_{idx:0>6d}.jpg'
            new_label_name = f'{data_category}_{idx:0>6d}.txt'

            if idx < train_img_cnt:
                new_image_path = os.path.join(res_path_dict['train_images_path'], f'{new_image_name}')
                new_label_path = os.path.join(res_path_dict['train_labels_path'], f'{new_label_name}')
            else:
                new_image_path = os.path.join(res_path_dict['valid_images_path'], f'{new_image_name}')
                new_label_path = os.path.join(res_path_dict['valid_labels_path'], f'{new_label_name}')

            # cp origin image to my dataset.
            shutil.copyfile(img_path, new_image_path)

            # write info to label
            with open(new_label_path, 'w') as file:
                file.write(label_info)

            # Use json fmt
            with open(save_info_path, 'a') as file:
                json.dump({'path': new_image_path, 'info': img_info_dict}, file)
                file.write('\n')


if __name__ == '__main__':
    data_process = DataProcess()
    data_process.get_all_ccpd_img_list('/home/dataset_b/car_license_data/CCPD2020', '/home/dataset_b/car_license_data/car_data')
    # YoloFmtDataset().get_car_license_info('0130387931034-90_87-237&454_447&522-447&526_245&527_237&459_439&458-0_0_2_15_24_24_30-123-72.jpg')
